import gym
import numpy as np
import time
# 定义系统参数
g = 9.81  # 重力加速度
m = 0.1   # 杆的质量
M = 1.0   # 小车的质量
l = 0.5   # 杆的长度

# 滑模面设计
def sliding_surface(state):
    x, dx, theta, dtheta = state
    s = dx + 5 * dtheta  # 滑模面设计
    return s

# 趋近律设计
def reaching_law(s, k=1.0):
    return -k * np.sign(s)

# 控制律设计
def control_law(state, s, u_reach):
    x, dx, theta, dtheta = state
    # 计算控制输入
    u = u_reach - (m * l * dtheta**2 * np.sin(theta) + m * g * np.cos(theta) * np.sin(theta)) / (M + m * np.cos(theta)**2)
    return u

# 初始化环境
env = gym.make("CartPole-v1", render_mode="human")
state, info = env.reset()

# 控制器运行
try:
    for t in range(100000):
        # 计算滑模面
        s = sliding_surface(state)
        
        # 计算趋近律
        u_reach = reaching_law(s)
        
        # 计算控制输入
        u = control_law(state, s, u_reach)
        
        # 将控制输入转换为动作 (0: 向左, 1: 向右)
        action = 0 if u > 0 else 1
        
        # 执行动作
        state, reward, done, truncated, info = env.step(action)
        
        # 渲染环境
        env.render()
        
        # 检查是否完成
        if done:
            print(f"Episode finished after {t + 1} timesteps")
            break
        
        # 添加一个小的延迟以便观察
        time.sleep(0.01)
except KeyboardInterrupt:
    print("Simulation stopped manually after {} timesteps".format(t))

env.close()